import librosa
import numpy as np
import os
import glob
from scipy import stats
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

def compute_kl(audio_path1, audio_path2):
    """
    计算两个音频文件之间的KL散度
    """
    try:
        # 加载音频
        audio1, sr1 = librosa.load(audio_path1)
        audio2, sr2 = librosa.load(audio_path2)
        
        # 确保采样率一致
        if sr1 != sr2:
            audio2 = librosa.resample(audio2, orig_sr=sr2, target_sr=sr1)
        
        # 提取MFCC特征
        mfcc1 = librosa.feature.mfcc(y=audio1, sr=sr1, n_mfcc=13)[0]  # 取第一维
        mfcc2 = librosa.feature.mfcc(y=audio2, sr=sr1, n_mfcc=13)[0]  # 使用相同的sr
        
        # 创建直方图分布
        bins = np.linspace(min(np.min(mfcc1), np.min(mfcc2)), 
                          max(np.max(mfcc1), np.max(mfcc2)), 100)
        hist1, _ = np.histogram(mfcc1, bins=bins, density=True)
        hist2, _ = np.histogram(mfcc2, bins=bins, density=True)
        
        # 归一化并平滑
        epsilon = 1e-10
        hist1 = hist1 + epsilon
        hist2 = hist2 + epsilon
        hist1 = hist1 / np.sum(hist1)
        hist2 = hist2 / np.sum(hist2)
        
        # 计算KL散度
        kl = np.sum(hist1 * np.log(hist1 / hist2))
        return kl
    except Exception as e:
        print(f"Error processing {audio_path1} and {audio_path2}: {e}")
        return None

def analyze_audio_files():
    """
    分析目录中的音频文件，计算KL散度并生成报告
    """
    # 定义原始音频文件
    fireworks_ori = "Fireworks._High-quality,_stereo.wav"  # 请确认文件名正确
    sports_car_ori = "Sports_car_passing_b.wav"
    
    # 获取当前目录下所有wav文件
    all_files = glob.glob("*.wav")
    
    # 分离出Fireworks和Sports_car相关文件
    fireworks_files = [f for f in all_files if "Fireworks" in f]
    sports_car_files = [f for f in all_files if "Sports_car" in f]
    
    # 确保原始文件存在
    if fireworks_ori not in all_files:
        print(f"Warning: {fireworks_ori} not found. Using first available Fireworks file as reference.")
        fireworks_ori = fireworks_files[0] if fireworks_files else None
    
    if sports_car_ori not in all_files:
        print(f"Warning: {sports_car_ori} not found. Using first available Sports_car file as reference.")
        sports_car_ori = sports_car_files[0] if sports_car_files else None
    
    # 存储结果
    results = []
    
    # 分析Fireworks文件
    if fireworks_ori:
        for file in fireworks_files:
            if file != fireworks_ori:
                kl = compute_kl(fireworks_ori, file)
                if kl is not None:
                    # 解析文件名获取参数信息
                    parts = file.replace(".wav", "").split("_")
                    inv_scale = None
                    rec_scale = None
                    steps = None
                    
                    for i, part in enumerate(parts):
                        if "invScale" in part:
                            inv_scale = float(part.replace("invScale", ""))
                        if "recScale" in part:
                            rec_scale = float(part.replace("recScale", ""))
                        if "steps" in part:
                            steps = int(part.replace("steps", ""))
                    
                    # 判断是revert还是ddim
                    method = "revert" if "revert" in file else "ddim"
                    
                    results.append({
                        "category": "Fireworks",
                        "method": method,
                        "inv_scale": inv_scale,
                        "rec_scale": rec_scale,
                        "steps": steps,
                        "file": file,
                        "kl_divergence": kl
                    })
    
    # 分析Sports_car文件
    if sports_car_ori:
        for file in sports_car_files:
            if file != sports_car_ori:
                kl = compute_kl(sports_car_ori, file)
                if kl is not None:
                    # 解析文件名获取参数信息
                    parts = file.replace(".wav", "").split("_")
                    inv_scale = None
                    rec_scale = None
                    steps = None
                    
                    for i, part in enumerate(parts):
                        if "invScale" in part:
                            inv_scale = float(part.replace("invScale", ""))
                        if "recScale" in part:
                            rec_scale = float(part.replace("recScale", ""))
                        if "steps" in part:
                            steps = int(part.replace("steps", ""))
                    
                    # 判断是revert还是ddim
                    method = "revert" if "revert" in file else "ddim"
                    
                    results.append({
                        "category": "Sports_car",
                        "method": method,
                        "inv_scale": inv_scale,
                        "rec_scale": rec_scale,
                        "steps": steps,
                        "file": file,
                        "kl_divergence": kl
                    })
    
    return results

def visualize_results(results):
    """
    可视化分析结果
    """
    if not results:
        print("No results to visualize")
        return
    
    # 转换为DataFrame
    df = pd.DataFrame(results)
    
    # 设置绘图风格
    sns.set(style="whitegrid")
    plt.figure(figsize=(15, 10))
    
    # 创建子图
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    
    # 1. 按类别和方法比较KL散度
    plt.subplot(2, 2, 1)
    sns.boxplot(x="category", y="kl_divergence", hue="method", data=df)
    plt.title("KL Divergence by Category and Method")
    plt.xticks(rotation=45)
    
    # 2. 按步数比较KL散度
    plt.subplot(2, 2, 2)
    sns.boxplot(x="steps", y="kl_divergence", hue="method", data=df)
    plt.title("KL Divergence by Steps and Method")
    
    # 3. 按inv_scale比较KL散度
    plt.subplot(2, 2, 3)
    sns.boxplot(x="inv_scale", y="kl_divergence", hue="method", data=df)
    plt.title("KL Divergence by Inv Scale and Method")
    
    # 4. 按rec_scale比较KL散度
    plt.subplot(2, 2, 4)
    sns.boxplot(x="rec_scale", y="kl_divergence", hue="method", data=df)
    plt.title("KL Divergence by Rec Scale and Method")
    
    plt.tight_layout()
    plt.savefig("kl_divergence_analysis.png")
    plt.show()
    
    # 创建详细的数据表格
    print("\n详细KL散度结果:")
    print(df.to_string(index=False))
    
    # 保存结果为CSV
    df.to_csv("kl_divergence_results.csv", index=False)
    print("\n结果已保存到 kl_divergence_results.csv")

# 执行分析
if __name__ == "__main__":
    print("开始分析音频文件KL散度...")
    results = analyze_audio_files()
    
    if results:
        print(f"成功分析了 {len(results)} 个文件")
        visualize_results(results)
    else:
        print("没有找到可分析的文件")